import logging
from http.server import HTTPServer, BaseHTTPRequestHandler

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
import json
import uuid
import time
from tools import tools
import datetime


class MCPServer:
    def __init__(self):
        self.protocol_version = "2025-06-18"
        self.server_info = {"name": "MCP服务器", "version": "1.0.0"}
        self.capabilities = {
            "tools": {},
            "resources": {},
            "prompts": {},
        }
        self.sessions = {}

    def _generate_session_id(self):
        return str(uuid.uuid4())

    def get_session_info(self, session_id):
        if session_id:
            return self.sessions[session_id]
        else:
            return None

    def _get_or_create_session(self, session_id):
        if not session_id:
            session_id = self._generate_session_id()
        if session_id not in self.sessions:
            self.sessions[session_id] = {
                "created_at": time.time(),
                "last_activity": time.time(),
                "message_count": 0,
                "client_info": None,
            }
            logger.info(f"创建新的会话:{session_id}")
        self.sessions[session_id]["last_activity"] = time.time()
        self.sessions[session_id]["message_count"] += 1
        return session_id

    def handle_initialize(self, params, session_id):
        # 获取客户端信息
        client_info = params.get("clientInfo", {})
        session_id = self._get_or_create_session(session_id)
        self.sessions[session_id]["client_info"] = client_info
        return {
            "protocolVersion": self.protocol_version,
            "capabilities": self.capabilities,
            "server_info": self.server_info,
            "session_id": session_id,
        }

    def handle_tools_call(self, params, session_id):
        if session_id:
            self._get_or_create_session(session_id)
        calls = params.get("calls", [])
        results = []
        for call in calls:
            tool_name = call.get("name")
            arguments = call.get("arguments", {})
            if tool_name == "calculate":
                expression = arguments.get("expression", "")
                result = eval(expression)
                results.append(
                    {
                        "name": tool_name,
                        "contents": [
                            {"type": "text", "text": f"计算结果:{expression}={result}"}
                        ],
                    }
                )
            else:
                results.append(
                    {
                        "name": tool_name,
                        "isError": True,
                        "error": {"message": f"未知工具:{tool_name}"},
                    }
                )
        return {"calls": results}

    def handle_tools_list(self, params, session_id):
        if session_id:
            self._get_or_create_session(session_id)
        return {"tools": tools}


class MCPHTTPHandler(BaseHTTPRequestHandler):
    def __init__(self, *args, mcp_server=None, **kwargs):
        self.mcp_server = mcp_server
        super().__init__(*args, **kwargs)

    def _handle_mcp_message(self, message, session_id):
        method = message.get("method")
        params = message.get("params")
        msg_id = message.get("id")
        if method == "initialize":
            result = self.mcp_server.handle_initialize(params, session_id)
            return {"id": msg_id, "result": result} if msg_id else None
        elif method == "tools/list":
            result = self.mcp_server.handle_tools_list(params, session_id)
            return {"id": msg_id, "result": result} if msg_id else None
        elif method == "tools/call":
            result = self.mcp_server.handle_tools_call(params, session_id)
            return {"id": msg_id, "result": result} if msg_id else None
        else:
            if msg_id:
                return {
                    "id": msg_id,
                    "error": {"code": -32601, "message": f"method not found:{method}"},
                }
            else:
                return None

    def _send_sse_event(self, event_type, data):
        try:
            event_data = f"event: {event_type}\n"
            event_data += f"data: {json.dumps(data)}\n\n"
            self.wfile.write(event_data.encode("utf-8"))
            self.wfile.flush()
        except Exception as e:
            logger.error(f"发送SSE事件错误:{str(e)}")

    def do_GET(self):
        try:
            if self.path != "/mcp":
                self.send_error(404, "MCP endpoint not found")
                return
            accept_header = self.headers.get("Accept", "")
            if "text/event-stream" not in accept_header:
                self.send_error(405, "Method not Allowed")
                return
            session_id = self.headers.get("Mcp-Session-Id")
            if session_id:
                session_info = self.mcp_server.get_session_info(session_id)
                if not session_info:
                    self.send_error(401, "Invalid session_id")
                    return
                else:
                    logger.info(f"SSE连接使用现有会话:{session_id}")
            else:
                session_id = self.mcp_server._get_or_create_session(None)
                logger.info("SSE连接创建新的会话:{session_id}")
            self.send_response(200)
            self.send_header("Content-Type", "text/event-stream")
            self.send_header("Cache-Control", "no-cache")
            self.send_header("Connection", "keep-alive")
            self.send_header("MCP-Protocol-Version", self.mcp_server.protocol_version)
            self.send_header("Mcp-Session-Id", session_id)
            self.end_headers()
            self._send_sse_event(
                "connected", {"status": "connected", "session_id": session_id}
            )
            logger.info(f"SSE连接已经建立，会话ID:{session_id}")
            try:
                while True:
                    time.sleep(5)
                    current_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                    self._send_sse_event(
                        "time_update",
                        {"server_time": current_time, "session_id": session_id},
                    )
                    logger.info(
                        f"发送时间更新:server_time:{current_time},session_id:{session_id}"
                    )
            except (BrokenPipeError, ConnectionResetError):
                logger.info(f"SSE连接已经断开，会话ID:{session_id}")
            except Exception as e:
                logger.info(f"SSE连接已经断开，会话ID:{session_id}")
        except Exception as e:
            logger.error(f"GET处理错误:{str(e)}")
            self.send_error(500, "Internal Server Error")

    def do_POST(self):
        try:
            if self.path != "/mcp":
                self.send_error(404, "MCP endpoint not found")
                return

            accept_header = self.headers.get("Accept", "")
            if (
                "application/json" not in accept_header
                and "text/event-stream" not in accept_header
            ):
                self.send_error(400, "Missing required Accept header")
                return
            protocol_version = self.headers.get("MCP-Protocol-Version", "")
            if (
                protocol_version
                and protocol_version != self.mcp_server.protocol_version
            ):
                self.send_error(
                    400, f"Unsupported protocol version :{protocol_version}"
                )
                return
            session_id = self.headers.get("Mcp-Session-Id")
            content_length = int(self.headers.get("Content-Length"), 0)
            body = self.rfile.read(content_length)
            try:
                message = json.loads(body.decode("utf-8"))
                response = self._handle_mcp_message(message, session_id)

                if response:
                    self.send_response(200)
                    self.send_header("Content-Type", "application/json")
                    self.send_header(
                        "MCP-Protocol-Version", self.mcp_server.protocol_version
                    )
                    self.end_headers()
                    self.wfile.write(json.dumps(response).encode("utf-8"))
                else:
                    self.send_response(202)
                    self.end_headers()
            except json.JSONDecodeError:
                self.send_error(400, "Invalid JSON Format")
                return

        except Exception as e:
            logger.error(f"POST处理错误:{str(e)}")
            self.send_error(500, "Internal Server Error")


def run_http_server(mcp_server, host, port):
    def handler_factory(*args, **kwargs):
        return MCPHTTPHandler(*args, mcp_server=mcp_server, **kwargs)

    server = HTTPServer((host, port), handler_factory)
    logger.info(f"HTTP服务器运行在 http://{host}:{port}/mcp")
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        logger.info("HTTP被中断,已经停止运行")
    finally:
        server.server_close()


def main():
    import argparse

    parser = argparse.ArgumentParser(description="MCP服务器")
    parser.add_argument(
        "--host", default="localhost", help="HTTP服务器主机:(默认是localhost)"
    )
    parser.add_argument(
        "--port", type=int, default=8000, help="HTTP服务器端口号:(默认是8000)"
    )
    args = parser.parse_args()
    mcp_server = MCPServer()
    run_http_server(mcp_server, args.host, args.port)


if __name__ == "__main__":
    main()
